import argparse
import scipy
from scipy import ndimage
import numpy as np
import sys
from packaging import version

import torch
from torch.autograd import Variable
import torchvision.models as models
import torch.nn.functional as F
from torch.utils import data, model_zoo
from model.deeplab import Res_Deeplab
from model.deeplab_multi import DeeplabMulti
from model.deeplab_vgg import DeeplabVGG
from dataset.vaihingen_dataset import VaihingenDataSet
from collections import OrderedDict
import os
from PIL import Image
import matplotlib.pyplot as plt
import torch.nn as nn

IMG_MEAN = np.array((97.6399, 92.5452, 85.9160), dtype=np.float32)   # Order: IR -> R -> G

DATA_DIRECTORY = '/home/allen/mm_dataset/vaihingen_600_normal'
DATA_LIST_PATH = './dataset/vaihingen_list/val.txt'
SAVE_PATH = '/media/allen/backup/毕业论文实验结果/potsdam_to_vaihingen_adaptseg/small_img'

IGNORE_LABEL = 255
REDUCE_ZERO_LABEL = True
NUM_CLASSES = 5
NUM_STEPS = 898 # Number of images in the validation set.
RESTORE_FROM = '/media/allen/backup/毕业论文实验结果/potsdam_to_vaihingen_adaptseg/potsdam_adaptseg_80k.pth'
# RESTORE_FROM_VGG = 'http://vllab.ucmerced.edu/ytsai/CVPR18/GTA2Cityscapes_vgg-ac4ac9f6.pth'
# RESTORE_FROM_ORC = 'http://vllab1.ucmerced.edu/~whung/adaptSeg/cityscapes_oracle-b7b9934.pth'
SET = 'val'

MODEL = 'DeeplabMulti'

# R, G, B values for each class to draw the prediction results.
# The order of class is ['build', 'car', 'grass', 'impervious', 'tree']
palette = [0, 0, 255, 255, 255, 0, 0, 255, 255, 255, 255, 255, 0, 255, 0]

# The len of palette should be 256 * 3, and the rest values are all 0(black).
zero_pad = 256 * 3 - len(palette)
for i in range(zero_pad):
    palette.append(0)


def colorize_mask(mask):
    # mask: numpy array of the mask
    new_mask = Image.fromarray(mask.astype(np.uint8)).convert('P')
    new_mask.putpalette(palette)

    return new_mask

def get_arguments():
    """Parse all the arguments provided from the CLI.

    Returns:
      A list of parsed arguments.
    """
    parser = argparse.ArgumentParser(description="DeepLab-ResNet Network")
    parser.add_argument("--model", type=str, default=MODEL,
                        help="Model Choice (DeeplabMulti/DeeplabVGG/Oracle).")
    parser.add_argument("--data-dir", type=str, default=DATA_DIRECTORY,
                        help="Path to the directory containing the Cityscapes dataset.")
    parser.add_argument("--data-list", type=str, default=DATA_LIST_PATH,
                        help="Path to the file listing the images in the dataset.")
    parser.add_argument("--ignore-label", type=int, default=IGNORE_LABEL,
                        help="The index of the label to ignore during the training.")
    parser.add_argument("--num-classes", type=int, default=NUM_CLASSES,
                        help="Number of classes to predict (including background).")
    parser.add_argument("--restore-from", type=str, default=RESTORE_FROM,
                        help="Where restore model parameters from.")
    parser.add_argument("--gpu", type=int, default=0,
                        help="choose gpu device.")
    parser.add_argument("--set", type=str, default=SET,
                        help="choose evaluation set.")
    parser.add_argument("--save", type=str, default=SAVE_PATH,
                        help="Path to save result.")
    return parser.parse_args()


def main():
    """Create the model and start the evaluation process."""

    args = get_arguments()

    gpu0 = args.gpu

    if not os.path.exists(args.save):
        os.makedirs(args.save)

    if args.model == 'DeeplabMulti':
        model = DeeplabMulti(num_classes=args.num_classes)

    if args.restore_from[:4] == 'http' :
        saved_state_dict = model_zoo.load_url(args.restore_from)
    else:
        saved_state_dict = torch.load(args.restore_from, map_location='cuda:0')
    ### for running different versions of pytorch
    # model_dict = model.state_dict()
    # saved_state_dict = {k: v for k, v in saved_state_dict.items() if k in model_dict}
    # model_dict.update(saved_state_dict)
    ###
    model.load_state_dict(saved_state_dict)

    model.eval()
    model.cuda(gpu0)

    testloader = data.DataLoader(VaihingenDataSet(args.data_dir, args.data_list, crop_size=(600, 600), mean=IMG_MEAN,
                                                  scale=False, mirror=False, set=args.set,
                                                  reduce_zero_label=REDUCE_ZERO_LABEL),
                                 batch_size=1, shuffle=False, pin_memory=True)

    interp = nn.Upsample(size=(600, 600), mode='bilinear', align_corners=True)

    for index, batch in enumerate(testloader):
        if index % 100 == 0:
            print('%d processd' % index)
        image, _, _, name = batch
        if args.model == 'DeeplabMulti':
            with torch.no_grad():
                output1, output2 = model(Variable(image).cuda(gpu0))
            output = interp(output2).cpu().data[0].numpy()

        output = output.transpose(1, 2, 0)
        output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)

        output_col = colorize_mask(output)
        output = Image.fromarray(output)

        name = name[0].split('/')[-1]
        output.save('%s/%s' % (args.save, name))
        output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0]))


if __name__ == '__main__':
    main()
